"""
MDC 工具函数和装饰器

提供方便的装饰器来自动管理函数级别的 MDC 上下文
"""

import functools
import inspect
import logging
from typing import Callable, Optional

from app.utils.mdc_context import mdc, mdc_context


def with_mdc(**mdc_kwargs):
    """
    为函数自动添加 MDC 上下文的装饰器

    用法：
    @with_mdc(component="user_service", operation="get_user")
    async def get_user(user_id: int):
        # 函数内的日志会自动包含 component 和 operation 上下文
        logger.info(f"获取用户信息: {user_id}")

    Args:
        **mdc_kwargs: 要设置的 MDC 上下文键值对
    """

    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs):
            with mdc_context(**mdc_kwargs):
                return await func(*args, **kwargs)

        @functools.wraps(func)
        def sync_wrapper(*args, **kwargs):
            with mdc_context(**mdc_kwargs):
                return func(*args, **kwargs)

        if inspect.iscoroutinefunction(func):
            return async_wrapper
        else:
            return sync_wrapper

    return decorator


def trace_function(operation: Optional[str] = None, component: Optional[str] = None):
    """
    函数追踪装饰器，自动记录函数执行时间和结果

    用法：
    @trace_function(operation="create_user", component="user_service")
    async def create_user(user_data: dict):
        # 函数执行会被自动追踪
        return await user_repo.create(user_data)

    Args:
        operation: 操作名称
        component: 组件名称
    """

    def decorator(func: Callable) -> Callable:
        func_name = operation or func.__name__
        comp_name = component or func.__module__

        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs):
            import time

            start_time = time.time()

            # 设置追踪上下文
            trace_context = {
                "function": func_name,
                "component": comp_name,
                "start_time": start_time,
            }

            with mdc_context(**trace_context):
                logger = logging.getLogger("mdc.trace")
                logger.debug(f"开始执行: {func_name}")

                try:
                    result = await func(*args, **kwargs)
                    execution_time = time.time() - start_time

                    mdc.put("execution_time_ms", round(execution_time * 1000, 2))
                    mdc.put("success", True)

                    logger.debug(f"执行成功: {func_name} - {execution_time:.3f}s")
                    return result

                except Exception as e:
                    execution_time = time.time() - start_time
                    mdc.put("execution_time_ms", round(execution_time * 1000, 2))
                    mdc.put("success", False)
                    mdc.put("error", str(e))

                    logger.error(f"执行失败: {func_name} - {execution_time:.3f}s - 错误: {str(e)}")
                    raise

        @functools.wraps(func)
        def sync_wrapper(*args, **kwargs):
            import time

            start_time = time.time()

            # 设置追踪上下文
            trace_context = {
                "function": func_name,
                "component": comp_name,
                "start_time": start_time,
            }

            with mdc_context(**trace_context):
                logger = logging.getLogger("mdc.trace")
                logger.debug(f"开始执行: {func_name}")

                try:
                    result = func(*args, **kwargs)
                    execution_time = time.time() - start_time

                    mdc.put("execution_time_ms", round(execution_time * 1000, 2))
                    mdc.put("success", True)

                    logger.debug(f"执行成功: {func_name} - {execution_time:.3f}s")
                    return result

                except Exception as e:
                    execution_time = time.time() - start_time
                    mdc.put("execution_time_ms", round(execution_time * 1000, 2))
                    mdc.put("success", False)
                    mdc.put("error", str(e))

                    logger.error(f"执行失败: {func_name} - {execution_time:.3f}s - 错误: {str(e)}")
                    raise

        if inspect.iscoroutinefunction(func):
            return async_wrapper
        else:
            return sync_wrapper

    return decorator


class MDCLogger:
    """MDC 增强的日志记录器"""

    def __init__(self, name: str = None):
        self.logger = logging.getLogger(name or __name__)

    def _enhance_message(self, message: str) -> str:
        """增强日志消息，包含 MDC 上下文摘要"""
        request_id = mdc.get("request_id")
        user_id = mdc.get("user_id")

        enhanced = message
        if request_id:
            enhanced = f"[{request_id}] {enhanced}"
        if user_id:
            enhanced = f"[用户:{user_id}] {enhanced}"

        return enhanced

    def debug(self, msg: str, *args, **kwargs):
        self.logger.debug(self._enhance_message(msg), *args, **kwargs)

    def info(self, msg: str, *args, **kwargs):
        self.logger.info(self._enhance_message(msg), *args, **kwargs)

    def warning(self, msg: str, *args, **kwargs):
        self.logger.warning(self._enhance_message(msg), *args, **kwargs)

    def error(self, msg: str, *args, **kwargs):
        self.logger.error(self._enhance_message(msg), *args, **kwargs)

    def critical(self, msg: str, *args, **kwargs):
        self.logger.critical(self._enhance_message(msg), *args, **kwargs)

    def exception(self, msg: str, *args, **kwargs):
        self.logger.exception(self._enhance_message(msg), *args, **kwargs)


# 创建全局 MDC 日志记录器
mdc_logger = MDCLogger("mdc")
